Performance Results:
  Dimensions: 32x1024x4096 -> 7168
  Triton: 5.17 ms (372.40 TFLOPS)
  PyTorch: 15.85 ms (121.38 TFLOPS)
  Speedup: 3.07x

Paper table format:
8       256     128     4096    7168    372 TFLOPS      5 ms    3.1x

Overall test result: All tests passed!

Verifying backward pass correctness...
Outputs match: True
Input gradients match: True
Weight gradients match: True

All gradients match! Running performance benchmark...
PyTorch backward time: 1.24 ms
Triton backward time: 0.25 ms
Speedup: 4.99x

contiguous grouped GEMM forward pass for MoE.
    All tokens mapped to the same expert must be in contiguous blocks of size group_size_m.

    Args:
        inputs: Input tensor of shape [M_total, K]
        expert_weights: Expert weight tensor of shape [num_experts, N, K]
        expert_indices: Indices tensor of shape [M_total] mapping each token to its expert
        group_size_m: Size of contiguous token blocks for each expert (default: 128)

    Returns:
        Output tensor of shape [M_total, N]